#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cublas_v2.h>
#include "utils.h"
#include <cuda_runtime.h>
#include "cublas_utils.h"

#include "check_point/rl0/mp_gemm.cpp"

int main(int argc, char *argv[]) {
    cudaSetDevice(0);
    std::vector<Case>  cases = {};

    #ifdef PERF
    float m_start = 100000;
    int loop = 10000;
    
    for(int m_id = 1*m_start; m_id <= 1*m_start; m_id+=2*m_start){
        cases.push_back({m_id, 1, 384, loop});
        for(int n_id = 8; n_id <= 128; n_id += 8){
            cases.push_back({m_id, n_id, 384, loop});
        }
    }
    #endif

    #ifdef EVAL
    const int m = atoi(argv[1]), n = atoi(argv[2]), k = atoi(argv[3]), loop = atoi(argv[4]);
    const std::string A_gt_path = argv[5], B_gt_path = argv[6], C_my_path = argv[7];
    cases.push_back({m, n, k, loop});
    #endif

    #ifdef PROFILE
    cases.push_back({262133, 64, 384, 1});
    #endif

    for(size_t case_id = 0; case_id < cases.size(); ++case_id){
        #ifdef PERF
        if(case_id > 0){
            cases[case_id].loop = max(100, int(1000.0f/cases[case_id-1].avg_latency))*3;
        }
        #endif

        Case c = cases[case_id];
        const int m = c.m, n = c.n, k = c.k, loop = c.loop;
        cudaStream_t stream = NULL;
        CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));

        uint8_t *d_A = nullptr;
        uint8_t *d_B = nullptr;
        int32_t *d_C = nullptr;
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_A), sizeof(uint8_t)*m*k/8));
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_B), sizeof(uint8_t)*k*n));
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_C), sizeof(int32_t)*n*m));

        #ifdef EVAL
        loadbin(d_A, m*k/8, A_gt_path);
        loadbin(d_B, sizeof(uint8_t)*k*n, B_gt_path);
        #endif
        {
            int start = now();
            for(int i = 0; i < loop; ++i){
                mp_gemm_w1a8(m, n, k, d_A, k/8, d_B, k, d_C, m, stream);
            }
            cudaStreamSynchronize(stream);
            int end = now();
            
            cases[case_id].avg_latency =float(end-start)/float(loop);

            #ifdef EVAL
            dumpbin(d_C, n*m*sizeof(int32_t), C_my_path);
            #endif
        }
        #ifdef PERF
        {
            Case c = cases[case_id];
            const int m = c.m, n = c.n, k = c.k, loop = c.loop;
            float avg_latency = c.avg_latency;
            printf("m %d n %d k %d loop %d avg_latency %0.4f ms  \n", m, n, k, loop, avg_latency);
        }
        #endif
        CUDA_CHECK(cudaFree(d_A));
        CUDA_CHECK(cudaFree(d_B));
        CUDA_CHECK(cudaFree(d_C));
        CUDA_CHECK(cudaStreamDestroy(stream));
    }
    return 0;
}